#include <vector>
#include <list>
#include <map>
#include <set>
#include <queue>
#include <deque>
#include <stack>
#include <bitset>
#include <algorithm>
#include <functional>
#include <numeric>
#include <utility>
#include <sstream>
#include <iostream>
#include <iomanip>
#include <cstdio>
#include <cmath>
#include <cstdlib>
#include <ctime>
#include <fstream>
#include <set>
#include "tree.h"

using namespace std;

bool isBalancedHelper(Tree *root, int &lh, int &rh)
{
    if (!root) {
	lh = rh = 0;
	return true;
    }
  
    int a = 0, b = 0;
    bool left = isBalancedHelper(root->left, a, b);
    lh = 1 + max(a,b);

    bool right = isBalancedHelper(root->right, a, b);
    rh = 1 + max(a,b);

    if (left && right && abs(lh - rh) <= 1) return true;
    return false;
}


bool isBalanced(Tree *root) {
    int lh, rh;
    bool ans = isBalancedHelper(root, lh, rh);
    if (ans && abs(lh - rh) <= 1) return true;
    return false; 
}


//add much better code
bool isBalancedHelper(TreeNode *root, int &h) 
{
    if (!root) {
        h = 0;
        return true;
    }   
    int lh, rh; 
    if (!isBalancedHelper(root->left, lh)) return false;
    if (!isBalancedHelper(root->right, rh)) return false;

    if (abs(lh - rh) > 1) return false;
    h = max(lh, rh) + 1;
    return true;
}

    bool isBalanced(TreeNode *root) {
        int h;
        return isBalancedHelper(root, h); 
    }   



int main(int argc, char **argv)
{
    int input[] = {21,7,5,6,3,2,1,4,6,23,43,54,12,54,8,10,11,35,56};
    //int input[] = {7};
    Tree *root = NULL, *root2 = NULL;
    int result;

    for (int i = 0; i < sizeof(input)/sizeof(int); i++) {
        TreeInsert(root, input[i]);
    }

    print_ascii_tree(root);
    cout << isBalanced(root) << endl;

}
